"""
意图识别模型的封装
"""
import config
from fasttext import FastText


class Classify:
    def __init__(self):
        """
        记载训练好的模型
        """
        self.model = FastText.load_model(config.classify_model_final_path)  # 词语特征模型
        self.model_by_word = FastText.load_model(config.classify_model_final_path_by_word)  # 单词特征模型

    def predict(self, sentence):
        """
        预测输入的数据结果，准确率
        :param sentence_cuted::分词之后的句子:{"cut_byword":str,"cut":str}
        :return:
        """
        result1 = self.model.predict(sentence["cut"])
        result2 = self.model_by_word.predict(sentence["cut_by_word"])
        for label, acc, label_by_word, acc_by_word in zip(*result1, *result2):
            # 把所有的label
            if label == "__label__chat":
                label = "__label__QA"
                acc = 1 - acc
            if label_by_word == "__label__chat":
                label_by_word = "__label__QA"
                acc_by_word = 1 - acc_by_word

            # 判断准确率
            if acc > 0.95 or acc_by_word > 0.95:  # 设置阈值
                return ("QA", max(acc, acc_by_word))
            else:
                return ("chat", 1 - min(acc - acc_by_word))

            # # TODO 假设有三个类别
            # if label == label_by_word:
            #     if acc > 0.95 or acc_by_word > 0.95:
            #         return label, max(acc, acc_by_word)
            #     else:
            #         return None, 0
            #
            # else:
            #     if acc_by_word > 0.99:
            #         return label_by_word, acc_by_word
            #     elif acc > 0.98:
            #         return label, acc
            #     else:
            #         return None, 0
